from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate,MessagesPlaceholder

from langgraph.prebuilt import create_react_agent
from typing import Annotated
from model_utils import getLLM
from app_utils import websearch
"""
特别说明：
如果没有查询出来：需要多次查询：
根据提供的信息，虽然我们没有具体的过去5年的GDP数据，但可以推测出在2019年至2023年间，英国的经济确实经历了一些挑战。
获取英国过去5年的国内生产总值和画一个折线图代码
1.构建web搜索工具
2.构建Researcher类
1)初始researcher  __init__ 预构建带工具的_agent，生成_chain
2)__call__方法   调用 _chain.invoke方法调用web工具， 获取英国过去5年的国内生产总值和画一个折线图代码
3.main函数调用测试
"""

@tool
def web_search(query:Annotated[str,"查询字段信息"]):
    '''查找互联网信息'''
    return websearch(query)

_research_system_template = '''
你是一个优秀的助手，根据提供工具回答问答
工具：{tools_name}
'''

class Researcher:
    def __init__(self,llm):
        _tools = [web_search]

        #消息站位
        _prompt = ChatPromptTemplate.from_messages(
            [
                ("system",_research_system_template),
                MessagesPlaceholder(variable_name="messages"),
            ]
        )

        _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools]))
        _agent = create_react_agent(llm,_tools)

        self._chain = _prompt | _agent

    def __call__(self,state):
        print("researcher node state",state)
        _rt = self._chain.invoke(state)
        _messages = _rt["messages"]
        return _messages[-1]

if __name__ == '__main__':
    from langchain_openai import ChatOpenAI
    _llm = getLLM()
    _researcher = Researcher(_llm)
    _query = "获取英国过去5年的国内生产总值，然后画一个折线图，一旦你把它编码好，并执行该代码，保存到当前路径下 名字为1.png。"
    #_query = "Fetch the UK's GDP over the past 5 years,  then draw a line graph of it.  Once you code it up, finish"
    _rt = _researcher({"messages": ["human", _query]})
    print(_rt)